#!/usr/bin/env python3

import os
import re
import sys
import argparse

DESCR = """afl-qemu-ts-cov1.0
Written and maintained by Andrea Fioraldi <andreafioraldi@gmail.com>
"""

opt = argparse.ArgumentParser(description=DESCR, formatter_class=argparse.RawTextHelpFormatter)
opt.add_argument('-i', help="Input directory", action='store', required=True)
opt.add_argument('-l', help="Log CSV of afl-qemu-cov", action='store', required=True)
opt.add_argument('-o', help="Output CSV file (default stdout)", action='store')

args = opt.parse_args()

if not os.path.isdir(args.i):
    print("ERROR: %s is not a directory" % args.i)
    exit(1)

aflpp_name_re = re.compile("id:(\d\d\d\d\d\d).*time:(\d+).+")
id_to_ts = {}

for testcase in os.listdir(args.i):
    if not testcase.startswith("id:"):
        continue
    r = aflpp_name_re.findall(testcase)
    if len(r) == 0:
        continue
    testcase_id = str(int(r[0][0]))
    testcase_ts = r[0][1]
    id_to_ts[testcase_id] = testcase_ts

with open(args.l) as f:
    log = filter(lambda x: not x.startswith("#"), f.readlines())
    log = map(lambda x: map(lambda y: y.strip(), x.split(",")), log)

new_cov = {}
for testcase_id, bb_addr in log:
    new_cov[testcase_id] = new_cov.get(testcase_id, 0)
    new_cov[testcase_id] += 1

if args.o:
    out = open(args.o, "w")
else:
    out = sys.stdout

out.write("# time, testcase_id, bbs\n")

total_cov = 0
for testcase_id in sorted(id_to_ts.keys(), key=lambda x: int(x)):
    if testcase_id not in new_cov:
        continue
    total_cov += new_cov[testcase_id]
    o = "%s, %s, %d\n" % (id_to_ts[testcase_id], testcase_id, total_cov)
    out.write(o)

out.close()
